package graph;

import java.util.*;

public class Prim_231115 {
    public static Set<Edge<Integer>> prim(Graph<Integer> graph) {
        if (graph.nodeMap.isEmpty()) {
            return Collections.emptySet();
        }

        Set<Edge<Integer>> result = new HashSet<>(graph.nodeMap.size());
        Set<Node<Integer>> nodeSet = new HashSet<>(graph.nodeMap.size());
        PriorityQueue<Edge<Integer>> pq = new PriorityQueue<>(Comparator.comparingInt(item -> item.weight));
        // 选择一个起始点
        for (Node<Integer> node : graph.nodeMap.values()) {
            if (nodeSet.contains(node)) {
                continue;
            }

            // Node<Integer> node = graph.nodeMap.values().stream().findFirst().get();
            nodeSet.add(node);
            // 初始化所有边
            node.nextEdges.forEach(pq::offer);

            // M 边数
            // O(M*logM)
            while (!pq.isEmpty() && nodeSet.size() != graph.nodeMap.size()) {
                Edge<Integer> edge = pq.poll();
                if (!nodeSet.contains(edge.to)) {
                    nodeSet.add(edge.to);
                    // 扩展出更多边，添加到小根堆中
                    edge.to.nextEdges.forEach(pq::offer);
                    result.add(edge);
                }
            }

            // break; // 如果图中存在森林，则需要遍历所有节点，没有森林则随便取一个点即可
        }

        return result;
    }
}
